package hw2;

import java.util.HashMap;
import java.util.Iterator;

public class Learner
{
	HashMap<Character, Integer> unigramModel;
	HashMap<Character, HashMap<Character, Integer>> bigramModel;
	HashMap<Character, HashMap<Character, HashMap<Character, Integer>>> trigramModel;
	int bestMethod;
	//trainingseen/traininglength = testseen/testlength
	HashMap<Character, Integer> trainingSeen;
	int trainingLength;
	HashMap<Character, Integer> testPredicted;
	HashMap<Character, Integer> testSeen;
	
	public Learner()
	{
		unigramModel = new HashMap<Character, Integer>();
		bigramModel = new HashMap<Character, HashMap<Character, Integer>>();
		trigramModel = new HashMap<Character, HashMap<Character, HashMap<Character, Integer>>>();
		bestMethod = 0;
		testPredicted = new HashMap<Character, Integer>();
		trainingSeen = new HashMap<Character, Integer>();
		testSeen = new HashMap<Character, Integer>();
	}
	
	public void train(String trainingInput)
	{
		int characterIndex = 0;
		char previousCharacter = 0;
		char previousPreviousCharacter = 0;
		trainingLength = trainingInput.length();
		while (characterIndex < trainingInput.length())
		{
			char currentCharacter = trainingInput.charAt(characterIndex);
			
			// Train the unigram model.
			if (unigramModel.get(currentCharacter) == null)
				unigramModel.put(currentCharacter, 0);
			
			unigramModel.put(currentCharacter, unigramModel.get(currentCharacter) + 1);
			
			// Add stuff to the proportion model.
			if (trainingSeen.get(currentCharacter) == null)
				trainingSeen.put(currentCharacter, 0);
			
			trainingSeen.put(currentCharacter, trainingSeen.get(currentCharacter) + 1);
			
			// Train the bigram model.
			if (bigramModel.get(previousCharacter) == null)
				bigramModel.put(previousCharacter, new HashMap<Character, Integer>());
			
			if (bigramModel.get(previousCharacter).get(currentCharacter) == null)
				bigramModel.get(previousCharacter).put(currentCharacter, 0);
			
			bigramModel.get(previousCharacter).put(currentCharacter, 
					bigramModel.get(previousCharacter).get(currentCharacter) + 1);
			
			// Train the trigram model.
			if (trigramModel.get(previousPreviousCharacter) == null)
				trigramModel.put(previousPreviousCharacter, new HashMap<Character, HashMap<Character, Integer>>());
			
			HashMap<Character, HashMap<Character, Integer>> innerHash = trigramModel.get(previousPreviousCharacter);
			if (innerHash.get(previousCharacter) == null)
				innerHash.put(previousCharacter, new HashMap<Character, Integer>());
			
			HashMap<Character, Integer> innerInnerHash = innerHash.get(previousCharacter);
			if (innerInnerHash.get(currentCharacter) == null)
				innerInnerHash.put(currentCharacter, 0);
			
			innerInnerHash.put(currentCharacter, innerInnerHash.get(currentCharacter) + 1);
			
			characterIndex++;
			previousPreviousCharacter = previousCharacter;
			previousCharacter = currentCharacter;
		}
		
		if (bigramModel.get(previousCharacter) == null)
			bigramModel.put(previousCharacter, new HashMap<Character, Integer>());
		
		if (trigramModel.get(previousCharacter) == null)
			trigramModel.put(previousCharacter, new HashMap<Character, HashMap<Character, Integer>>());
	}
	
	public double unigramTest(String testInput, boolean smart)
	{
		int numCorrectPredictions = 0;

		//Replace this with actual test prediction HACK
		int characterIndex = 0;
		if(testPredicted.isEmpty()){
			Iterator<Character> trainingSeenIterator = trainingSeen.keySet().iterator(); 
			while (trainingSeenIterator.hasNext()) {
				char nextCharacter = trainingSeenIterator.next();
				testPredicted.put(nextCharacter,trainingSeen.get(nextCharacter)*testInput.length()/trainingLength);
			}
		}
		
		while (characterIndex < testInput.length())
		{
			char predictedCharacter;
			// Predict next character.
			if (smart){
				predictedCharacter = unigramPrediction2();
			}
			else {
				predictedCharacter = unigramPrediction();
			}
			
			// Read character.
			char actualNextCharacter = testInput.charAt(characterIndex);
			
			// Determine if prediction was correct.
			if (predictedCharacter == actualNextCharacter)
				numCorrectPredictions++;
			
			// Update model.
			if (unigramModel.get(actualNextCharacter) == null)
				unigramModel.put(actualNextCharacter, 0);
			if(testSeen.get(actualNextCharacter) == null)
				testSeen.put(actualNextCharacter, 0);
			
			unigramModel.put(actualNextCharacter, unigramModel.get(actualNextCharacter) + 1);
			testSeen.put(actualNextCharacter, testSeen.get(actualNextCharacter) + 1);
			
			characterIndex++;
		}
		
		// Print percentage of correct predictions.
		
		return (numCorrectPredictions / (float) testInput.length()) * 100 ; 
		
		
		
	}
	
	public void unigramTestRandom() {
		// Print random prediction performance (1 / number of unique symbols).
		System.out.println("Random prediction performance: " + 
			(1 / (float) unigramModel.keySet().size()) * 100 + "%");
	}
	
	public double bigramTest(String testInput, boolean smart)
	{
		int numCorrectPredictions = 0;

		//Replace this with actual test prediction HACK
		if(testPredicted.isEmpty()){
			Iterator<Character> trainingSeenIterator = trainingSeen.keySet().iterator(); 
			while (trainingSeenIterator.hasNext()) {
				char nextCharacter = trainingSeenIterator.next();
				testPredicted.put(nextCharacter,trainingSeen.get(nextCharacter)*testInput.length()/trainingLength);
			}
		}
		
		int characterIndex = 0;
		char previousCharacter = 0;
		while (characterIndex < testInput.length())
		{
			// Predict next character.
			char predictedCharacter = bigramPrediction(previousCharacter);
			
			// Read character.
			char actualNextCharacter = testInput.charAt(characterIndex);
			
			// Determine if prediction was correct.
			if (predictedCharacter == actualNextCharacter)
				numCorrectPredictions++;
			
			// Update unigram model.
			if (unigramModel.get(actualNextCharacter) == null)
				unigramModel.put(actualNextCharacter, 0);
			
			unigramModel.put(actualNextCharacter, unigramModel.get(actualNextCharacter) + 1);
			
			// Update bigram model.
			if (bigramModel.get(actualNextCharacter) == null)
				bigramModel.put(actualNextCharacter, new HashMap<Character, Integer>());
			
			if (bigramModel.get(previousCharacter).get(actualNextCharacter) == null)
				bigramModel.get(previousCharacter).put(actualNextCharacter, 0);
			
			bigramModel.get(previousCharacter).put(actualNextCharacter, 
					bigramModel.get(previousCharacter).get(actualNextCharacter) + 1);
			
			characterIndex++;
			previousCharacter = actualNextCharacter;
		}
		
		// Print percentage of correct predictions.
		return (numCorrectPredictions / (float) testInput.length()) * 100; 
		
		
	}
	public void bigramTestRandom() {
		// Print random prediction performance (1 / number of unique symbols).
		System.out.println("Random prediction performance: " + 
				(1 / (float) (bigramModel.keySet().size() - 1)) * 100 + "%");
	}
	public double trigramTest(String testInput, boolean smart)
	{
		int numCorrectPredictions = 0;

		//Replace this with actual test prediction HACK
		if(testPredicted.isEmpty()){
			Iterator<Character> trainingSeenIterator = trainingSeen.keySet().iterator(); 
			while (trainingSeenIterator.hasNext()) {
				char nextCharacter = trainingSeenIterator.next();
				testPredicted.put(nextCharacter,trainingSeen.get(nextCharacter)*testInput.length()/trainingLength);
			}
		}
		
		int characterIndex = 0;
		char previousCharacter = 0;
		char previousPreviousCharacter = 0;
		while (characterIndex < testInput.length())
		{
			// Predict next character.
			char predictedCharacter = trigramPrediction(previousCharacter, previousPreviousCharacter);
			
			// Read character.
			char actualNextCharacter = testInput.charAt(characterIndex);
			
			// Determine if prediction was correct.
			if (predictedCharacter == actualNextCharacter)
				numCorrectPredictions++;
			
			// Update unigram model.
			if (unigramModel.get(actualNextCharacter) == null)
				unigramModel.put(actualNextCharacter, 0);
			
			unigramModel.put(actualNextCharacter, unigramModel.get(actualNextCharacter) + 1);
			
			// Update bigram model.
			if (bigramModel.get(actualNextCharacter) == null)
				bigramModel.put(actualNextCharacter, new HashMap<Character, Integer>());
			
			if (bigramModel.get(previousCharacter).get(actualNextCharacter) == null)
				bigramModel.get(previousCharacter).put(actualNextCharacter, 0);
			
			bigramModel.get(previousCharacter).put(actualNextCharacter, 
					bigramModel.get(previousCharacter).get(actualNextCharacter) + 1);
			
			// Update trigram model.
			if (trigramModel.get(actualNextCharacter) == null)
				trigramModel.put(actualNextCharacter, new HashMap<Character, HashMap<Character, Integer>>());
			
			HashMap<Character, HashMap<Character, Integer>> innerHash = trigramModel.get(previousPreviousCharacter);
			if (innerHash.get(previousCharacter) == null)
				innerHash.put(previousCharacter, new HashMap<Character, Integer>());
			
			HashMap<Character, Integer> innerInnerHash = innerHash.get(previousCharacter);
			if (innerInnerHash.get(actualNextCharacter) == null)
				innerInnerHash.put(actualNextCharacter, 0);
			
			innerInnerHash.put(actualNextCharacter, innerInnerHash.get(actualNextCharacter) + 1);
			
			characterIndex++;
			previousPreviousCharacter = previousCharacter;
			previousCharacter = actualNextCharacter;
		}
		
		// Print percentage of correct predictions.
		return (numCorrectPredictions / (float) testInput.length()) * 100; 
		
		
	}
	
	public void trigramTestRandom() {
		// Print random prediction performance (1 / number of unique symbols).
		System.out.println("Random prediction performance: " + 
		(1 / (float) (trigramModel.keySet().size() - 1)) * 100 + "%");
	}
	
	private char unigramPrediction()
	{
		int highestFrequency = -1;
		char mostFrequentCharacter = 0;
		
		Iterator<Character> modelIterator = unigramModel.keySet().iterator();
		while (modelIterator.hasNext())
		{
			char nextCharacter = modelIterator.next();
			int nextFrequency = unigramModel.get(nextCharacter); 
			if (nextFrequency > highestFrequency)
			{
				highestFrequency = nextFrequency;
				mostFrequentCharacter = nextCharacter;
			}
		}
		//System.out.println("Unigram prediction: " + mostFrequentCharacter + ", frequency " + highestFrequency);
		return mostFrequentCharacter;
	}
	
	private char unigramPrediction2()
	{
		int highestFrequency = -1;
		char mostFrequentCharacter = 0;
		//trainingseen/traininglength = testseen/testlength
		Iterator<Character> modelIterator = unigramModel.keySet().iterator();
		while (modelIterator.hasNext())
		{
			char nextCharacter = modelIterator.next();
			int nextFrequency = unigramModel.get(nextCharacter); 
			
			//probabilistically change next character's likelihood
			if(testSeen.get(nextCharacter) == null)
				testSeen.put(nextCharacter, 0);
			if(testPredicted.get(nextCharacter) != null)
				nextFrequency = nextFrequency * (1 - (testSeen.get(nextCharacter)/testPredicted.get(nextCharacter)));
			
			if (nextFrequency > highestFrequency)
			{
				highestFrequency = nextFrequency;
				mostFrequentCharacter = nextCharacter;
			}
		}
		//System.out.println("Unigram prediction: " + mostFrequentCharacter + ", frequency " + highestFrequency);
		return mostFrequentCharacter;
	}
	
	private char bigramPrediction(char previousCharacter)
	{
		int highestFrequency = -1;
		char mostFrequentCharacter = 0;
		
		if (bigramModel.get(previousCharacter) == null || 
				bigramModel.get(previousCharacter).size() == 0)
			return unigramPrediction();
		
		Iterator<Character> modelIterator = 
			bigramModel.get(previousCharacter).keySet().iterator();
		
		while (modelIterator.hasNext())
		{
			char nextCharacter = modelIterator.next();
			int nextFrequency = bigramModel.get(previousCharacter).get(nextCharacter);
			if (nextFrequency > highestFrequency)
			{
				highestFrequency = nextFrequency;
				mostFrequentCharacter = nextCharacter;
			}
		}
		//System.out.println("Bigram prediction: " + mostFrequentCharacter + ", frequency " + highestFrequency);
		return mostFrequentCharacter;
	}
	
	private char trigramPrediction(char previousCharacter, char previousPreviousCharacter)
	{
		int highestFrequency = -1;
		char mostFrequentCharacter = 0;
		
		// If the trigram model can't make a prediction, give the bigram model a try.
		if (trigramModel.get(previousPreviousCharacter) == null || 
				trigramModel.get(previousPreviousCharacter).size() == 0)
			return bigramPrediction(previousCharacter);
		
		HashMap<Character, HashMap<Character, Integer>> innerHash = trigramModel.get(previousPreviousCharacter);
		if (innerHash.get(previousCharacter) == null || 
				innerHash.get(previousCharacter).size() == 0)
			return bigramPrediction(previousCharacter);
		
		HashMap<Character, Integer> innerInnerHash = innerHash.get(previousCharacter);
		Iterator<Character> modelIterator = innerInnerHash.keySet().iterator();
		
		while (modelIterator.hasNext())
		{
			char nextCharacter = modelIterator.next();
			int nextFrequency = innerInnerHash.get(nextCharacter);
			if (nextFrequency > highestFrequency)
			{
				highestFrequency = nextFrequency;
				mostFrequentCharacter = nextCharacter;
			}
		}
		//System.out.println("Trigram prediction: " + mostFrequentCharacter + ", frequency " + highestFrequency);
		return mostFrequentCharacter;
	}
	
	public void setBestMethod(int bestMethod)
	{
		this.bestMethod = bestMethod;
	}
	
	public int getBestMethod()
	{
		return bestMethod;
	}
}
